package com.bins.langchain.langchain4j.service.impl;

import com.bins.langchain.langchain4j.service.ChatService;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.client.MilvusServiceClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;

@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    private OllamaChatModel ollamaChatModel;

    @Autowired
    private OllamaStreamingChatModel ollamaStreamingChatModel;

    @Autowired
    private OllamaEmbeddingModel ollamaEmbeddingModel;

    @Autowired
    private MilvusEmbeddingStore milvusEmbeddingStore;

    @Autowired
    private MilvusServiceClient milvusServiceClient;

    private final String TEXTID = "text_id";
    private final String VECTOR = "vector";

    @Override
    public String generate(String prompt) {
        return ollamaChatModel.generate(prompt);
    }

    @Override
    public String embedding(String text) {
        Response<Embedding> embed = ollamaEmbeddingModel.embed(text);
        Embedding embedding = embed.content();
        return milvusEmbeddingStore.add(embedding);
    }

    @Override
    public String embeddingWithMeta(String text) {
        return null;
    }

    @Override
    public List<String> search(String text) {
        return null;
    }

    @Override
    public Flux<String> streamChat(String message) {
        return Flux.create(emitter -> {
            ollamaStreamingChatModel.generate(message, new StreamingResponseHandler<AiMessage>() {
                @Override
                public void onNext(String token) {
                    emitter.next(token);
                }

                @Override
                public void onComplete(Response<AiMessage> response) {
                    emitter.complete();
                }

                @Override
                public void onError(Throwable error) {
                    emitter.error(error);
                }
            });
        });
    }
}
